import pickle
from argparse import ArgumentParser
from ibm_model import IBMModel1


def argmax(a):
    if len(a) == 0:
        return None

    max_idx = 0
    for i in range(1, len(a)):
        if a[i] > a[max_idx]:
            max_idx = i
    return max_idx


def main(option):
    src_vocab = pickle.load(open(option.source_vocab, 'rb'))
    tgt_vocab = pickle.load(open(option.target_vocab, 'rb'))

    lf = len(src_vocab['word2id'])
    le = len(tgt_vocab['word2id'])

    model = pickle.load(open(option.pretrained_model, 'rb'))
    print('Loaded model ' + option.pretrained_model)

    samples = [line.strip() for line in open(option.input_file, 'r')]
    print('Loaded input file ' + option.input_file)

    with open(option.output_file, 'w') as f:
        for src_word in samples:
            src_id = src_vocab['word2id'][src_word]
            probs = [model.t[i][src_id] for i in range(le)]
            tgt_id = argmax(probs)
            tgt_word = tgt_vocab['id2word'][tgt_id]
            f.write(tgt_word + '\n')
    print('Output saved to ' + option.output_file)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--source_vocab', type=str, default='output/fbis.en.pkl')
    parser.add_argument('--target_vocab', type=str, default='output/fbis.zh.pkl')
    parser.add_argument('--pretrained_model', type=str, default='output/model.e10.pkl')
    parser.add_argument('--input_file', type=str, default='output/input.txt')
    parser.add_argument('--output_file', type=str, default='output/output.txt')
    option = parser.parse_args()

    main(option)
